Best time to buy and sell stock IV

Time: O(KxN); Space: O(K); hard

Say you have an array for which the i-th element is the price of a given stock on day i.

Design an algorithm to find the maximum profit. You may complete at most k transactions.

Note:

  • You may not engage in multiple transactions at the same time (ie, you must sell the stock before you buy again).

Example 1:

Input: prices = [2,4,1], k = 2

Output: 2

Explanation:

  • Buy on day 1 (price = 2) and sell on day 2 (price = 4), profit = 4-2 = 2.

Example 2:

Input: prices = [3,2,6,5,0,3], k = 2

Output: 7

Explanation:

  • Buy on day 2 (price = 2) and sell on day 3 (price = 6), profit = 6-2 = 4.

  • Then buy on day 5 (price = 0) and sell on day 6 (price = 3), profit = 3-0 = 3.

Example 3:

Input: prices = [4, 4, 6, 1, 1, 4, 2 ,5], k = 2

Output: 6

Explanation:

  • Buy at 4 and sell at 6. Then buy at 1 and sell at 5. Your profit is 2 + 4 = 6.

Follow up:

  • O(KxN) time. N is the size of prices.

[6]:

class Solution1(object):
    """
    Time: O(K*N)
    Space: O(K)
    """
    def maxProfit(self, prices, k):
        """
        :type prices: List[int]
        :type k: int
        :rtype int
        """
        if k >= len(prices) // 2:
            return self.maxAtMostNPairsProfit(prices)

        return self.maxAtMostKPairsProfit(prices, k)

    def maxAtMostNPairsProfit(self, prices):
        profit = 0
        for i in range(len(prices) - 1):
            profit += max(0, prices[i + 1] - prices[i])
        return profit

    def maxAtMostKPairsProfit(self, prices, k):
        max_buy = [float("-inf") for _ in range(k + 1)]
        max_sell = [0 for _ in range(k + 1)]

        for i in range(len(prices)):
            for j in range(1, min(k, i//2 + 1) + 1):
                max_buy[j] = max(max_buy[j], max_sell[j-1] - prices[i])
                max_sell[j] = max(max_sell[j], max_buy[j] + prices[i])

        return max_sell[k]

[5]:

s = Solution1()

prices = [2,4,1]
k = 2
assert s.maxProfit(prices, k) == 2

prices = [3,2,6,5,0,3]
k = 2
assert s.maxProfit(prices, k) == 7

prices = [4, 4, 6, 1, 1, 4, 2 ,5]
k = 2
assert s.maxProfit(prices, k) == 6